use smallvec::SmallVec;
use std::collections::BTreeMap;

pub struct TrieMap<T> {
    children: Vec<(char, TrieNode<T>)>,
}

fn root_key<T>(&(c, _): &(char, TrieNode<T>)) -> char {
    c
}

struct TrieNode<T> {
    children: BTreeMap<char, TrieNode<T>>,
    values: SmallVec<[T; 1]>,
}

impl<T: std::fmt::Debug> TrieMap<T> {
    pub fn new() -> Self {
        Self {
            children: Vec::new(),
        }
    }
    pub(super) fn push_root_char_values(&mut self, key: char, values: impl IntoIterator<Item = T>) {
        self.children.push((key, TrieNode::from_values(values)));
    }
    pub(super) fn sort_root(&mut self) {
        self.children.sort_by_key(root_key);
        self.children.dedup_by(|(c1, node1), (c2, node2)| {
            if c1 == c2 {
                node2.values.append(&mut node1.values);
                return true;
            }
            false
        });
    }
    fn insert_root_char(&mut self, key: char) -> &mut TrieNode<T> {
        match self.children.binary_search_by_key(&key, root_key) {
            Ok(i) => &mut self.children[i].1,
            Err(i) => {
                self.children.insert(i, (key, TrieNode::new()));
                &mut self.children[i].1
            }
        }
    }
    pub fn insert_all(
        &mut self,
        key: impl IntoIterator<Item = char>,
        values: impl IntoIterator<Item = T>,
    ) {
        let mut key = key.into_iter();
        let Some(first) = key.next() else { return };
        self.insert_root_char(first).extend(key, values);
    }
    pub fn insert(&mut self, key: impl IntoIterator<Item = char>, value: T) {
        let mut key = key.into_iter();
        let Some(first) = key.next() else { return };
        self.insert_root_char(first).push(key, value);
    }
    pub fn get<'a, 'b>(&'a self, key: &'b str) -> Option<(&'a [T], &'b str)> {
        let mut chars = key.chars();
        let i = self
            .children
            .binary_search_by_key(&chars.next()?, root_key)
            .ok()?;
        self.children[i].1.get(chars.as_str())
    }
    pub fn from_root_chars(keys: impl Iterator<Item = (char, impl Iterator<Item = T>)>) -> Self {
        let children: Vec<_> = keys
            .map(|(c, values)| (c, TrieNode::from_values(values)))
            .collect();
        let mut this = Self { children };
        this.sort_root();
        this
    }
}

impl<T> TrieNode<T> {
    fn new() -> Self {
        Self {
            children: BTreeMap::new(),
            values: SmallVec::new(),
        }
    }
    fn from_values(values: impl IntoIterator<Item = T>) -> Self {
        Self {
            children: BTreeMap::new(),
            values: values.into_iter().collect(),
        }
    }
    fn entry(&mut self, key: impl IntoIterator<Item = char>) -> &mut Self {
        let mut node = self;
        for c in key {
            node = node.children.entry(c).or_default();
        }
        node
    }
    fn push(&mut self, key: impl IntoIterator<Item = char>, value: T) {
        self.entry(key).values.push(value);
    }
    fn extend(&mut self, key: impl IntoIterator<Item = char>, values: impl IntoIterator<Item = T>) {
        self.entry(key).values.extend(values);
    }
    fn non_empty_value_and_remainder<'a, 'b>(
        &'a self,
        remainder: &'b str,
    ) -> Option<(&'a [T], &'b str)> {
        (!self.values.is_empty()).then_some((&self.values, remainder))
    }
    fn get<'a, 'b>(&'a self, key: &'b str) -> Option<(&'a [T], &'b str)> {
        self.get_child(key)
            .or_else(|| self.non_empty_value_and_remainder(key))
    }
    fn get_child<'a, 'b>(&'a self, key: &'b str) -> Option<(&'a [T], &'b str)> {
        let mut chars = key.chars();
        self.children.get(&chars.next()?)?.get(chars.as_str())
    }
}

impl<K, V, T: std::fmt::Debug> Extend<(K, V)> for TrieMap<T>
where
    K: IntoIterator<Item = char>,
    V: IntoIterator<Item = T>,
{
    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
        for (key, values) in iter {
            self.insert_all(key, values);
        }
    }
}

impl<K, V, T: std::fmt::Debug> FromIterator<(K, V)> for TrieMap<T>
where
    K: IntoIterator<Item = char>,
    V: IntoIterator<Item = T>,
{
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        let mut trie = TrieMap::new();
        trie.extend(iter);
        trie
    }
}

impl<T: std::fmt::Debug> Default for TrieMap<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Default for TrieNode<T> {
    fn default() -> Self {
        Self::new()
    }
}
